import os
from idlelib.iomenu import encoding
from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union, Type
from uuid import UUID
from langchain.utilities import SerpAPIWrapper
from langchain.agents import Tool
from langchain.tools.file_management.write import WriteFileInput
from langchain.tools.file_management.read import ReadFileTool
from langchain_core.callbacks import StdOutCallbackHandler, FileCallbackHandler
from langchain_openai import OpenAIEmbeddings
import faiss
from langchain.vectorstores import FAISS
from langchain.docstore import InMemoryDocstore
from langchain_experimental.autonomous_agents import AutoGPT
from langchain_openai import ChatOpenAI
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from langchain_community.tools.file_management.utils import (
    INVALID_PATH_TEMPLATE,
    BaseFileToolMixin,
    FileValidationError,
)

class CustomWriteFileTool(BaseFileToolMixin, BaseTool):
    """Tool that writes a file to disk."""

    name: str = "write_file"
    args_schema: Type[BaseModel] = WriteFileInput
    description: str = "Write file to disk"
    result:str = ''

    def _run(
        self,
        file_path: str,
        text: str,
        append: bool = False,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        self.result = text
        try:
            write_path = self.get_relative_path(file_path)

        except FileValidationError:
            return INVALID_PATH_TEMPLATE.format(arg_name="file_path", value=file_path)
        try:
            write_path.parent.mkdir(exist_ok=True, parents=False)
            mode = "a" if append else "w"
            with write_path.open(mode, encoding="utf-8") as f:
                f.write(text)

            # with self.get_relative_path('out.txt').open(mode, encoding='utf-8') as f:
            #     f.write(text)

            return f"File written successfully to {file_path}."
        except Exception as e:
            return "Error: " + str(e)

    # TODO: Add aiofiles method

# https://serpapi.com/
def search_from_www(serpapi_key, openai_key, openai_model, question):
    os.environ["SERPAPI_API_KEY"] = serpapi_key
    os.environ["OPENAI_API_KEY"] = openai_key
    # 构造 AutoGPT 的工具集
    search = SerpAPIWrapper()
    writeTool=CustomWriteFileTool()
    tools = [
        Tool(
            name="search",
            func=search.run,
            description="useful for when you need to answer questions about current events. You should ask targeted questions",
        ),
        writeTool,
        ReadFileTool(),
    ]

    # OpenAI Embedding 模型
    embeddings_model = OpenAIEmbeddings()
    # OpenAI Embedding 向量维数
    embedding_size = 1536
    # 使用 Faiss 的 IndexFlatL2 索引
    index = faiss.IndexFlatL2(embedding_size)
    # 实例化 Faiss 向量数据库
    vectorstore = FAISS(embeddings_model.embed_query, index, InMemoryDocstore({}), {})
    agent = AutoGPT.from_llm_and_tools(
        ai_name="Jarvis",
        ai_role="Assistant",
        tools=tools,
        llm=ChatOpenAI(model_name=openai_model, temperature=0, verbose=True),
        memory=vectorstore.as_retriever(
            search_type="similarity_score_threshold",
            search_kwargs={"score_threshold": 0.8}),  # 实例化 Faiss 的 VectorStoreRetriever
    )

    # 打印 Auto-GPT 内部的 chain 日志
    agent.chain.verbose = True
    agent.run([question])
    return writeTool.result



